"""
只专注一件事：修复错误的SQL
"""

from nl2sql.model.llm import LLM


class RepairAgent:
    def __init__(self, llm: LLM):
        self.llm = llm
        self.repair_prompt = None

    def repair(self, wrong_sql: str, error_info: str) -> str:
        prompt = self.repair_prompt.format(wrong_sql=wrong_sql, error_info=error_info)
        fixed_sql: str = self.llm.chat(prompt=prompt)
        return fixed_sql
